{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import pickle\n",
    "import numpy as np\n",
    "import os\n",
    "import json\n",
    "import random\n",
    "    \n",
    "# load source words\n",
    "source_words_path = os.path.join(os.getcwd(), 'source_words.pkl')\n",
    "with open(source_words_path, 'rb') as f_source_words:\n",
    "    source_words = pickle.load(f_source_words)\n",
    "    \n",
    "# load target words\n",
    "target_words_path = os.path.join(os.getcwd(), 'target_words.pkl')\n",
    "with open(target_words_path, 'rb') as f_target_words:\n",
    "    target_words = pickle.load(f_target_words)\n",
    "    \n",
    "# load label words\n",
    "label_words_path = os.path.join(os.getcwd(), 'label_words.pkl')\n",
    "with open(label_words_path, 'rb') as f_label_words:\n",
    "    label_words = pickle.load(f_label_words)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "945\n",
      "133\n",
      "27\n",
      "1\n",
      "3\n",
      "2\n",
      "0\n",
      "1\n"
     ]
    }
   ],
   "source": [
    "print(len(source_words))\n",
    "print(len(target_words))\n",
    "print(len(label_words))\n",
    "print(source_words['<pad>'])\n",
    "print(source_words['<eos>'])\n",
    "print(source_words['<sos>'])\n",
    "print(source_words['<unk>'])\n",
    "print(target_words['<pad>'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "编码器Encoder的实现\n",
    "'''\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "class CNNAttention(nn.Module):\n",
    "    def __init__(self, input_dim, intent_out, slot_out, hid_dim, n_layers, kernel_size, dropout, src_pad_idx, n_heads, max_length=50):\n",
    "        super(CNNAttention, self).__init__()\n",
    "        for kernel in kernel_size:\n",
    "            assert kernel % 2 == 1,'kernel size must be odd!' # 卷积核size为奇数，方便序列两边pad处理\n",
    "        \n",
    "        self.src_pad_idx = src_pad_idx\n",
    "        \n",
    "        self.scale = torch.sqrt(torch.FloatTensor([0.5])).to(device) # 确保整个网络的方差不会发生显著变化\n",
    "        \n",
    "        self.tok_embedding = nn.Embedding(input_dim, hid_dim) # token编码\n",
    "        self.pos_embedding = nn.Embedding(max_length, hid_dim) # token的位置编码\n",
    "        \n",
    "        self.hid2hid = nn.Linear(hid_dim * 2, hid_dim) # 线性层，从2 * hid_dim转为hid_dim\n",
    "        \n",
    "        # 不同的kernel_size\n",
    "        '''\n",
    "        self.conv_module = list()\n",
    "        for k in kernel_size:\n",
    "            conv = nn.ModuleList([nn.Conv1d(in_channels=hid_dim,\n",
    "                                                  out_channels=2*hid_dim, # 卷积后输出的维度，这里2*hid_dim是为了后面的glu激活函数\n",
    "                                                  kernel_size=k,\n",
    "                                                  padding=(k - 1)//2) # 序列两边补0个数，保持维度不变\n",
    "                                                  for _ in range(n_layers)])\n",
    "            self.conv_module.append(conv)\n",
    "        '''\n",
    "        \n",
    "        self.conv_1 = nn.ModuleList([nn.Conv1d(in_channels=hid_dim,\n",
    "                                                  out_channels=2*hid_dim, # 卷积后输出的维度，这里2*hid_dim是为了后面的glu激活函数\n",
    "                                                  kernel_size=kernel_size[0],\n",
    "                                                  padding=(kernel_size[0] - 1)//2) # 序列两边补0个数，保持维度不变\n",
    "                                                  for _ in range(n_layers)])\n",
    "        self.conv_2 = nn.ModuleList([nn.Conv1d(in_channels=hid_dim,\n",
    "                                                  out_channels=2*hid_dim, # 卷积后输出的维度，这里2*hid_dim是为了后面的glu激活函数\n",
    "                                                  kernel_size=kernel_size[1],\n",
    "                                                  padding=(kernel_size[1] - 1)//2) # 序列两边补0个数，保持维度不变\n",
    "                                                  for _ in range(n_layers)])\n",
    "        self.conv_3 = nn.ModuleList([nn.Conv1d(in_channels=hid_dim,\n",
    "                                                  out_channels=2*hid_dim, # 卷积后输出的维度，这里2*hid_dim是为了后面的glu激活函数\n",
    "                                                  kernel_size=kernel_size[2],\n",
    "                                                  padding=(kernel_size[2] - 1)//2) # 序列两边补0个数，保持维度不变\n",
    "                                                  for _ in range(n_layers)])\n",
    "        \n",
    "        # 几个卷积模块转换维度\n",
    "        self.convhid2hid = nn.Linear(len(kernel_size) * hid_dim, hid_dim)\n",
    "        \n",
    "        # 多头注意力模块\n",
    "        self.self_attention = MultiHeadAttentionLayer(hid_dim, n_heads, dropout)\n",
    "        \n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "        # intent detection 意图识别\n",
    "        self.intent_output = nn.Linear(hid_dim, intent_out)\n",
    "        \n",
    "         # slot filling，槽填充\n",
    "        self.slot_out = nn.Linear(hid_dim, slot_out)\n",
    "    \n",
    "    def make_src_mask(self, src):\n",
    "        # src: [batch_size, src_len]\n",
    "        src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2) # [batch_size, 1, 1, src_len]\n",
    "        \n",
    "        return src_mask\n",
    "        \n",
    "    def forward(self, src):\n",
    "        # src: [batch_size, src_len]\n",
    "        # src_mask: [batch_size, src_len]\n",
    "        batch_size = src.shape[0]\n",
    "        src_len = src.shape[1]\n",
    "        \n",
    "        src_mask = self.make_src_mask(src) # [batch_size, 1, 1, src_len]\n",
    "        \n",
    "        # 创建token位置信息\n",
    "        pos = torch.arange(src_len).unsqueeze(0).repeat(batch_size, 1).to(device) # [batch_size, src_len]\n",
    "        \n",
    "        # 对token与其位置进行编码\n",
    "        tok_embedded = self.tok_embedding(src) # [batch_size, src_len, hid_dim]\n",
    "        pos_embedded = self.pos_embedding(pos.long()) # [batch_size, src_len, hid_dim]\n",
    "        \n",
    "        # 对token embedded和pos_embedded逐元素加和\n",
    "        embedded = self.dropout(tok_embedded + pos_embedded) # [batch_size, src_len, hid_dim]\n",
    "        \n",
    "        # 转变维度，卷积在输入数据的最后一维进行\n",
    "        conv_input = embedded.permute(0, 2, 1) # [batch_size, hid_dim, src_len]\n",
    "        \n",
    "        '''\n",
    "        combine_conv_module_list = []\n",
    "        for conv_module in self.conv_module:\n",
    "            conved_input = conv_input\n",
    "            # 以下进行卷积块\n",
    "            for i, conv in enumerate(conv_module):\n",
    "                \n",
    "                # 进行卷积\n",
    "                conved = conv(self.dropout(conved_input)) # [batch_size, 2*hid_dim, src_len]\n",
    "\n",
    "                # 进行激活glu\n",
    "                conved = F.glu(conved, dim=1) # [batch_size, hid_dim, src_len]\n",
    "\n",
    "                # 进行残差连接\n",
    "                conved = (conved + conved_input) * self.scale # [batch_size, hid_dim, src_len]\n",
    "\n",
    "                # 作为下一个卷积块的输入\n",
    "                conved_input = conved\n",
    "                \n",
    "            combine_conv_module_list.append(conved)\n",
    "            \n",
    "        # 拼接几个卷积块特征: [batch_size, len(kernel_size) * hid_dim, src_len]\n",
    "        combine_conv_module = combine_conv_module_list[0]\n",
    "        for i in range(1, len(combine_conv_module_list)):\n",
    "            combine_conv_module = torch.cat([combine_conv_module, combine_conv_module_list[i]], dim = 1)\n",
    "        '''\n",
    "        \n",
    "        # 第一个kernel_size\n",
    "        conved_input = conv_input\n",
    "        for i, conv in enumerate(self.conv_1):\n",
    "            # 进行卷积\n",
    "            conved1 = conv(self.dropout(conved_input)) # [batch_size, 2*hid_dim, src_len]\n",
    "\n",
    "            # 进行激活glu\n",
    "            conved1 = F.glu(conved1, dim=1) # [batch_size, hid_dim, src_len]\n",
    "\n",
    "            # 进行残差连接\n",
    "            conved1 = (conved1 + conved_input) * self.scale # [batch_size, hid_dim, src_len]\n",
    "\n",
    "            # 作为下一个卷积块的输入\n",
    "            conved_input = conved1\n",
    "        \n",
    "        combine_conv_module = conved1\n",
    "        \n",
    "        # 第二个kernel_size\n",
    "        conved_input = conv_input\n",
    "        for i, conv in enumerate(self.conv_2):\n",
    "            # 进行卷积\n",
    "            conved2 = conv(self.dropout(conved_input)) # [batch_size, 2*hid_dim, src_len]\n",
    "\n",
    "            # 进行激活glu\n",
    "            conved2 = F.glu(conved2, dim=1) # [batch_size, hid_dim, src_len]\n",
    "\n",
    "            # 进行残差连接\n",
    "            conved2 = (conved2 + conved_input) * self.scale # [batch_size, hid_dim, src_len]\n",
    "\n",
    "            # 作为下一个卷积块的输入\n",
    "            conved_input = conved2\n",
    "            \n",
    "        combine_conv_module = torch.cat([combine_conv_module, conved2], dim = 1)\n",
    "        \n",
    "        # 第三个kernel_size\n",
    "        conved_input = conv_input\n",
    "        for i, conv in enumerate(self.conv_3):\n",
    "            # 进行卷积\n",
    "            conved3 = conv(self.dropout(conved_input)) # [batch_size, 2*hid_dim, src_len]\n",
    "\n",
    "            # 进行激活glu\n",
    "            conved3 = F.glu(conved3, dim=1) # [batch_size, hid_dim, src_len]\n",
    "\n",
    "            # 进行残差连接\n",
    "            conved3 = (conved3 + conved_input) * self.scale # [batch_size, hid_dim, src_len]\n",
    "\n",
    "            # 作为下一个卷积块的输入\n",
    "            conved_input = conved3\n",
    "            \n",
    "        combine_conv_module = torch.cat([combine_conv_module, conved3], dim = 1)\n",
    "        \n",
    "        \n",
    "        \n",
    "        conved = self.convhid2hid(combine_conv_module.permute(0, 2, 1)) # [batch_size, src_len, hid_dim]\n",
    "        \n",
    "        # 这里在所有卷积之后增加了一个多头自注意力层，它的输入是\n",
    "        self_attention, _ = self.self_attention(embedded, embedded, embedded, src_mask) # [batch_size, query_len, hid_dim]\n",
    "        \n",
    "        # 拼接卷积后的特征与多头注意力后的特征\n",
    "        combined_conv_attention = torch.cat([conved, self_attention], dim=2) # [batch_size, query_len, 2*hid_dim]\n",
    "        \n",
    "        # 经过一线性层，将2*hid_dim转为hid_dim，作为输出的特征\n",
    "        conved = self.hid2hid(combined_conv_attention) # [batch_size, query_len, hid_dim]\n",
    "        \n",
    "        # 又是一个残差连接，逐元素加和输出，作为encoder的联合输出特征\n",
    "        combined = (conved + embedded) * self.scale # [batch_size, src_len, hid_dim]\n",
    "        \n",
    "        # 意图识别,加一个平均池化,池化后的维度是：[batch_size, hid_dim]\n",
    "        intent_output = self.intent_output(self.dropout(F.max_pool1d(combined.permute(0, 2, 1), combined.shape[1]).squeeze())) # [batch_size, intent_dim]\n",
    "    \n",
    "        # 槽填充\n",
    "        slot_output = self.slot_out(self.dropout(combined)) # [batch_size, trg_len, output_dim]\n",
    "        \n",
    "        return intent_output, slot_output\n",
    " \n",
    "'''\n",
    "多头注意力multi-head attention\n",
    "'''\n",
    "class MultiHeadAttentionLayer(nn.Module):\n",
    "    def __init__(self, hid_dim, n_heads, dropout):\n",
    "        super(MultiHeadAttentionLayer, self).__init__()\n",
    "        \n",
    "        assert hid_dim % n_heads == 0\n",
    "        \n",
    "        self.hid_dim = hid_dim\n",
    "        self.n_heads = n_heads\n",
    "        self.head_dim = hid_dim // n_heads\n",
    "        \n",
    "        self.fc_q = nn.Linear(hid_dim, hid_dim)\n",
    "        self.fc_k = nn.Linear(hid_dim, hid_dim)\n",
    "        self.fc_v = nn.Linear(hid_dim, hid_dim)\n",
    "        \n",
    "        self.fc_o = nn.Linear(hid_dim, hid_dim)\n",
    "        \n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "        self.scale = torch.sqrt(torch.FloatTensor([self.hid_dim])).to(device) # 缩放因子\n",
    "        \n",
    "    def forward(self, query, key, value, mask=None):\n",
    "        '''\n",
    "        query: [batch_size, query_len, hid_dim]\n",
    "        key: [batch_size, key_len, hid_dim]\n",
    "        value: [batch_size, value_len, hid_dim]\n",
    "        '''\n",
    "        batch_size = query.shape[0]\n",
    "        \n",
    "        Q = self.fc_q(query) # [batch_size, query_len, hid_dim]\n",
    "        K = self.fc_k(key) # [batch_size, key_len, hid_dim]\n",
    "        V = self.fc_v(value) # [batch_size, value_len, hid_dim]\n",
    "        \n",
    "        Q = Q.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3) # [batch_size, n_heads, query_len, head_dim]\n",
    "        K = K.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3) # [batch_size, n_heads, key_len, head_dim]\n",
    "        V = V.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3) # [batch_size, n_heads, value_len, head_dim]\n",
    "        \n",
    "        # [batch_size, n_heads, query_len, head_dim] * [batch_size, n_heads, head_dim, key_len]\n",
    "        energy = torch.matmul(Q, K.permute(0, 1, 3, 2)) / self.scale # [batch_size, n_heads, query_len, key_len]\n",
    "        \n",
    "        if mask != None:\n",
    "            energy = energy.masked_fill(mask == 0, -1e10)\n",
    "        \n",
    "        attention = torch.softmax(energy, dim=-1) # [batch_size, n_heads, query_len, key_len]\n",
    "        \n",
    "        # [batch_size, n_heads, query_len, key_len] * [batch_size, n_heads, value_len, head_dim]\n",
    "        x = torch.matmul(self.dropout(attention), V) # [batch_size, n_heads, query_len, head_dim]\n",
    "        \n",
    "        x = x.permute(0, 2, 1, 3).contiguous() # [batch_size, query_len, n_heads, head_dim]\n",
    "        \n",
    "        x = x.view(batch_size, -1, self.hid_dim) # [batch_size, query_len, hid_dim]\n",
    "        \n",
    "        x = self.fc_o(x) # [batch_size, query_len, hid_dim]\n",
    "        \n",
    "        return x, attention\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hid_dim = 64\n",
    "conv_layers = 8\n",
    "kernel_size = (3,5,7)\n",
    "dropout = 0.5\n",
    "n_heads = 8\n",
    "\n",
    "model_path = os.path.join(os.getcwd(), \"model.h5\")\n",
    "\n",
    "input_dim = len(source_words) # source 词典大小（即词数量）\n",
    "slot_out = len(target_words) # target 词典大小（即实体类型数量）\n",
    "intent_out = len(label_words) # label 词典大小（即意图类别数量）\n",
    "\n",
    "src_pad_idx = source_words['<pad>']\n",
    "\n",
    "model = CNNAttention(input_dim, intent_out, slot_out, hid_dim, conv_layers, kernel_size, dropout, src_pad_idx, n_heads)\n",
    "\n",
    "model = model.to(device)\n",
    "\n",
    "model.load_state_dict(torch.load(model_path))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['<sos>', 'meals', 'are', 'served', 'on', 'american', 'flight', '665', '673', 'from', 'milwaukee', 'to', 'seattle', '<eos>']\n",
      "[2, 261, 30, 264, 8, 71, 11, 746, 747, 5, 65, 4, 110, 3]\n",
      "len source :14\n",
      "slot_prediciton:o o o o o b-airline_name o o o o b-fromloc.city_name o b-toloc.city_name o\n",
      "intent_prediction:flight\n"
     ]
    }
   ],
   "source": [
    "sentence = \"i would like to find a flight from charlotte to las vegas that makes a stop in st. louis\"\n",
    "sentence2 = \"which airlines have first class flights today\"\n",
    "sentence3 = 'on april first i need a ticket from tacoma to san jose departing before 7 am'\n",
    "sentence4 = 'what are the departure times from detroit to westchester county'\n",
    "sentence5 = 'please find a flight round trip from los angeles to tacoma washington with a stopover in san francisco not exceeding the price of 300 dollars for june tenth 1993'\n",
    "sentence6 = 'show me flight us 1207 from indianapolis to charlotte on monday and flight us 1500 from charlotte to minneapolis on monday and flight twa 639 from minneapolis to indianapolis'\n",
    "sentence7 = 'i need to fly from denver to westchester county on june seventh after 3 pm'\n",
    "sentence8 = 'what meals are served on american flight 811 from tampa to milwaukee'\n",
    "sentence9 = 'meals are served on american flight 665 673 from milwaukee to seattle'\n",
    "\n",
    "\n",
    "with torch.no_grad():\n",
    "    tokenized = sentence9.split()  # tokenize the sentence\n",
    "    tokenized = ['<sos>'] + tokenized + ['<eos>']\n",
    "    indexed = [source_words[t] for t in tokenized]  # convert to integer sequence\n",
    "    #pad = [1]*(seq_len - len(indexed))\n",
    "    #indexed.extend(pad)\n",
    "    print(tokenized)\n",
    "    print(indexed)\n",
    "    print('len source :{}'.format(len(indexed)))\n",
    "    src_tensor = torch.LongTensor(indexed)  # convert to tensor\n",
    "    src_tensor = src_tensor.unsqueeze(0).to(device)  # reshape in form of batch,no. of words\n",
    "    \n",
    "    intent_output, slot_output = model(src_tensor)  # prediction\n",
    "    \n",
    "    intent_output = intent_output.squeeze()\n",
    "    intent_output = intent_output.argmax()\n",
    "    intent = intent_output.detach().item()\n",
    "    \n",
    "    pred_token = slot_output.squeeze().argmax(1)\n",
    "    \n",
    "    slot_prediction = [target_words.itos[t.item()] for t in pred_token]\n",
    "    \n",
    "    print('slot_prediciton:{}'.format(' '.join(slot_prediction)))\n",
    "    print('intent_prediction:{}'.format(label_words.itos[intent]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "O O O B-airline_code B-flight_number O B-fromloc.city_name O B-toloc.city_name O B-depart_date.day_name O O O B-flight_number O B-fromloc.city_name O B-toloc.city_name O B-depart_date.day_name O O B-airline_code B-flight_number O B-fromloc.city_name O B-toloc.city_name"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
